#
# Stochastic Optimization PS#1 Problem 6a
# Nurse Staffing
#

# Reference Format: Vehicle Routing Problem

#
# Imports
#
from coopr.pyomo import *
#from coopr.opt.base import solver

#
# Model
#
model = AbstractModel()

#
# Parameters
#

# Define sets
model.I = Set() # units
model.J = Set()  # days

# Data_deterministic
model.s = Param(model.I)  # # of patients can be served each shift for each inital assignment
model.r = Param(model.I)  # # of patients can be served after each reassignment

model.Cost = Param()    # cost per shift
model.Penalty = Param() # penalty cost per unserved patient

# Data_stochastic
model.d = Param(model.I, model.J)

#
# Variables
#

# integer
model.X = Var(model.I, model.J, within=PositiveIntegers) # # of assigned nuerses to ith unit on jth day
# continuous
model.U = Var(model.I, model.J, within=PositiveReals) # # of reassigned nuerses from ith unit on jth day
model.V = Var(model.I, model.J, within=PositiveReals) # # of reassigned nerses to ith unit on jth day
model.W = Var(model.I, model.J, within=PositiveReals) # # of unserved patients on ith unit on jth day

model.FirstStageCost = Var()
model.SecondStageCost = Var()

#
# Stage-specific cost computations
#
def first_stage_cost_rule(mod):
    return (mod.FirstStageCost - mod.Cost * sum( mod.X[i,j] for i in mod.I for j in mod.J ) ) == 0.0
model.ComputeFirstStageCost = Constraint(rule=first_stage_cost_rule)


def second_stage_cost_rule(mod):
    expcost = mod.Penalty * sum( mod.W[i,j] for i in mod.I for j in mod.J )
    return (mod.SecondStageCost - expcost) == 0.0
model.ComputeSecondStageCost = Constraint(rule=second_stage_cost_rule)


#
# Constraints
#
def demand_rule(mod, i, j):
    return  mod.W[i,j] >= ( mod.d[i,j] - (mod.X[i,j] - mod.U[i,j]) * mod.s[i] - mod.V[i,j] * mod.r[i] )
model.DemandRule = Constraint(model.I, model.J, rule=demand_rule)

def reassign_dominate_rule(mod, i, j):
    return mod.X[i,j] >= mod.U[i,j]
model.ReassignDominateRule = Constraint(model.I, model.J, rule=reassign_dominate_rule)

def reassign_balance_rule(mod, j):
    return sum(mod.U[i,j] for i in mod.I) == sum(mod.V[i,j] for i in mod.I)
model.ReassignBalanceRule = Constraint(model.J, rule=reassign_balance_rule)
#
# Objective
#
def total_cost_rule(mod):
    return (mod.FirstStageCost + mod.SecondStageCost)


model.TotalCost = Objective(rule=total_cost_rule, sense=minimize)

#model.create('ReferenceModel.py')
